import numpy as np

class Bandit:
    def __init__(self,arms=10):
        self.rates=np.random.rand(arms)

    def play(self,arm):
        rate=self.rates[arm]
        if rate>np.random.rand():
            return 1
        else:
            return 0

# bandit=Bandit()
# for i in range(5):
#     print(bandit.play(0))

class Agent:
    def __init__(self,epsilon, action_size=10):
        # 初始化参数
        self.action_size = action_size  # 老虎机臂数
        self.epsilon = epsilon  # ε-贪婪策略的探索概率
        self.Qs = np.zeros(action_size)  # 每个臂的价值估计
        self.ns = np.zeros(action_size)  # 每个臂被选择的次数
    
    def get_action(self):
        # 实现ε-贪婪策略选择动作
        if np.random.random() < self.epsilon:
            # 探索：随机选择一个臂
            return np.random.randint(self.action_size)
        else:
            # 利用：选择当前价值估计最高的臂
            return np.argmax(self.Qs)
    
    def update(self, action, reward):
        # 更新选中臂的价值估计
        self.ns[action] += 1
        # 增量更新方式计算平均收益
        self.Qs[action] += (reward - self.Qs[action]) / self.ns[action]

# 使用示例
if __name__ == "__main__":
    bandit = Bandit(arms=10)
    agent = Agent(action_size=10, epsilon=0.1)
    
    # 进行1000次尝试
    total_reward = 0
    for t in range(1000):
        # 选择动作
        action = agent.get_action()
        # 获得奖励
        reward = bandit.play(action)
        # 更新价值估计
        agent.update(action, reward)
        total_reward += reward
    
    print(f"总收益: {total_reward}")
